"""
PDF文档内容提取工具

主要功能：
1. 提取PDF中的表格和文本内容
2. 智能识别段落边界
3. 支持跨页表格合并
4. 输出txt格式文件

核心函数：
- extract_content_all(): 主要提取函数，支持跨页表格合并
- extract_mixed_content(): 处理混合内容页面（表格+文本）
- clean_text(): 清理文本中的Unicode问题字符
- clean_table_data(): 清理表格数据
"""

import os
import pdfplumber
import re
import unicodedata

# 清理文本中的问题字符
def clean_text(text):
    """清理文本中的问题字符"""
    if not text:
        return text
    
    # 规范化Unicode字符
    text = unicodedata.normalize('NFKC', text)
    
    # 移除或替换常见的问题字符
    # 替换不可见字符和控制字符（保留换行符、制表符和回车符）
    text = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]', '', text)
    
    # 替换多个连续空格为单个空格
    text = re.sub(r'\s+', ' ', text)
    
    # 移除行首行尾空格
    text = text.strip()
    
    return text

# 清理表格数据中的问题字符
def clean_table_data(table_data):
    """清理表格数据中的问题字符"""
    if not table_data:
        return table_data
    
    cleaned_table = []
    for row in table_data:
        if row:
            cleaned_row = []
            for cell in row:
                if cell:
                    cleaned_cell = clean_text(str(cell))
                    cleaned_row.append(cleaned_cell)
                else:
                    cleaned_row.append(cell)
            cleaned_table.append(cleaned_row)
        else:
            cleaned_table.append(row)
    
    return cleaned_table


def extract_text_excluding_tables(page, table_bboxes):
    """
    从页面中提取文本，排除表格区域
    """
    paragraphs = []
    
    try:
        # 获取页面中的所有文本对象（字符级别）
        chars = page.chars
        if not chars:
            return paragraphs
        
        # 过滤掉在表格区域内的字符
        filtered_chars = []
        for char in chars:
            char_in_table = False
            char_x = (char['x0'] + char['x1']) / 2  # 字符中心x坐标
            char_y = (char['top'] + char['bottom']) / 2  # 字符中心y坐标
            
            # 检查字符是否在任何表格区域内
            for bbox in table_bboxes:
                x0, y0, x1, y1 = bbox
                # 添加一些边距以确保完全排除表格区域
                margin = 5
                if (x0 - margin <= char_x <= x1 + margin and 
                    y0 - margin <= char_y <= y1 + margin):
                    char_in_table = True
                    break
            
            if not char_in_table:
                filtered_chars.append(char)
        
        if not filtered_chars:
            return paragraphs
        
        # 按位置排序字符（从上到下，从左到右）
        chars_sorted = sorted(filtered_chars, key=lambda c: (-c['top'], c['x0']))
        
        # 将字符组织成行
        lines = []
        current_line = []
        last_top = None
        line_tolerance = 2  # 同一行的y坐标容差
        
        for char in chars_sorted:
            if last_top is None or abs(char['top'] - last_top) <= line_tolerance:
                current_line.append(char)
            else:
                if current_line:
                    # 按x坐标排序当前行的字符
                    current_line.sort(key=lambda c: c['x0'])
                    line_text = ''.join([c['text'] for c in current_line]).strip()
                    if line_text:
                        lines.append({
                            'text': line_text,
                            'top': current_line[0]['top'],
                            'bottom': max(c['bottom'] for c in current_line),
                            'x0': min(c['x0'] for c in current_line),
                            'x1': max(c['x1'] for c in current_line),
                            'chars': current_line
                        })
                current_line = [char]
            last_top = char['top']
        
        # 处理最后一行
        if current_line:
            current_line.sort(key=lambda c: c['x0'])
            line_text = ''.join([c['text'] for c in current_line]).strip()
            if line_text:
                lines.append({
                    'text': line_text,
                    'top': current_line[0]['top'],
                    'bottom': max(c['bottom'] for c in current_line),
                    'x0': min(c['x0'] for c in current_line),
                    'x1': max(c['x1'] for c in current_line),
                    'chars': current_line
                })
        
        # 将行组织成段落
        current_paragraph = []
        
        for i, line in enumerate(lines):
            if not current_paragraph:
                current_paragraph.append(line)
            else:
                # 计算行间距
                prev_line = current_paragraph[-1]
                line_spacing = prev_line['bottom'] - line['top']
                
                # 获取平均字符高度作为参考
                avg_char_height = sum(c['size'] for c in prev_line['chars']) / len(prev_line['chars'])
                
                # 段落分割条件
                is_new_paragraph = (
                    line_spacing > avg_char_height * 1.5 or  # 行间距较大
                    (line['x0'] - prev_line['x0'] > avg_char_height) or  # 明显缩进
                    prev_line['text'].rstrip().endswith(('.', '。', '!', '！', '?', '？', ':', '：')) or  # 句子结束
                    (len(prev_line['text']) < 50 and line_spacing > avg_char_height * 1.2)  # 短行后的间距
                )
                
                if is_new_paragraph:
                    # 完成当前段落
                    if current_paragraph:
                        paragraph_text = ' '.join([l['text'] for l in current_paragraph])
                        if paragraph_text.strip():
                            paragraphs.append({
                                "text": clean_text(paragraph_text),
                                "bbox": {
                                    "x0": min(l['x0'] for l in current_paragraph),
                                    "y0": min(l['top'] for l in current_paragraph),
                                    "x1": max(l['x1'] for l in current_paragraph),
                                    "y1": max(l['bottom'] for l in current_paragraph)
                                },
                                "line_count": len(current_paragraph),
                                "excluded_tables": True  # 标记这是排除表格后的文本
                            })
                    current_paragraph = [line]
                else:
                    current_paragraph.append(line)
        
        # 处理最后一个段落
        if current_paragraph:
            paragraph_text = ' '.join([l['text'] for l in current_paragraph])
            if paragraph_text.strip():
                paragraphs.append({
                    "text": clean_text(paragraph_text),
                    "bbox": {
                        "x0": min(l['x0'] for l in current_paragraph),
                        "y0": min(l['top'] for l in current_paragraph),
                        "x1": max(l['x1'] for l in current_paragraph),
                        "y1": max(l['bottom'] for l in current_paragraph)
                    },
                    "line_count": len(current_paragraph),
                    "excluded_tables": True  # 标记这是排除表格后的文本
                })
    
    except Exception as e:
        print(f"警告: 排除表格区域的文本提取失败: {e}")
        # 如果失败，尝试基本的文本提取（不排除表格区域）
        try:
            full_text = page.extract_text()
            if full_text:
                paragraphs.append({
                    "text": clean_text(full_text),
                    "bbox": None,
                    "line_count": 1,
                    "excluded_tables": False
                })
        except Exception as e2:
            print(f"警告: 基本文本提取也失败: {e2}")
    
    return paragraphs

# 提取页面中的混合内容（表格和文本）
def extract_mixed_content(page):
    """提取页面中的混合内容（表格和文本）- 优化版本，避免表格区域文本重复"""
    content_data = {
        "tables": [],
        "text_blocks": [],
        "full_text": "",
        "table_bboxes": []  # 新增：存储表格边界框信息
    }

    # draw_edges_text_blocks(page, "edges")

    # 首先提取所有表格和它们的边界框
    table_objects = page.find_tables(table_settings={
        "vertical_strategy":   "lines",   # 有竖线，继续用 lines
        "horizontal_strategy": "lines",   # 有横线
        "snap_tolerance":      12,        # 断线吸合（↑关键）
        "join_tolerance":      12,        # 断头延长（↑关键）
        "edge_min_length":     2,        # 别把细线滤掉
        "min_words_vertical":  0,        # 允许单元格只有 1 个字也认
        "min_words_horizontal":0,
    })
    
    table_bboxes = []
    if table_objects:
        print(f"找到 {len(table_objects)} 个表格对象")
        for i, table_obj in enumerate(table_objects):
            # 获取表格数据
            table_data = table_obj.extract()
            if table_data and len(table_data) > 0:
                cleaned_table = clean_table_data(table_data)
                
                # 获取表格边界框
                bbox = table_obj.bbox
                table_bboxes.append(bbox)
                
                content_data["tables"].append({
                    "table_index": i + 1,
                    "data": cleaned_table,
                    "row_count": len(table_data),
                    "col_count": len(table_data[0]) if table_data and table_data[0] else 0,
                    "bbox": {
                        "x0": bbox[0],
                        "y0": bbox[1], 
                        "x1": bbox[2],
                        "y1": bbox[3]
                    }
                })
        

    # 提取完整的页面文本（仅用于统计，不保存到最终数据）
    full_text = page.extract_text()
    content_data["_full_text_for_stats"] = clean_text(full_text) if full_text else ""
    
    # 提取非表格区域的文本
    non_table_paragraphs = extract_text_excluding_tables(page, table_bboxes)
    content_data["text_blocks"] = non_table_paragraphs
    
    # 绘制表格边界框用于调试（可选）
    # if table_bboxes:
    #     draw_table_bboxes(page, table_bboxes, "table_debug")
    
    return content_data


def is_page_number_format(text):
    """
    检查文本是否是页码格式（如 1-1-16）
    """
    if not text or not isinstance(text, str):
        return False
    
    # 匹配格式如 "1-1-16", "2-3-45" 等
    pattern = r'^\d+-\d+-\d+$'
    return bool(re.match(pattern, text.strip()))

def is_table_at_page_top(table_bbox, page_height, threshold=100):
    """
    判断表格是否在页面顶部
    Args:
        table_bbox: 表格边界框信息 {"x0": x0, "y0": y0, "x1": x1, "y1": y1}
        page_height: 页面高度
        threshold: 距离页面顶部的阈值（像素）
    Returns:
        bool: 如果表格在页面顶部返回True
    """
    # PDF坐标系中，y0越小越靠近页面顶部
    distance_from_top = table_bbox["y0"]
    return distance_from_top <= threshold

def is_table_at_page_bottom(table_bbox, page_height, threshold=100):
    """
    判断表格是否在页面底部
    Args:
        table_bbox: 表格边界框信息 {"x0": x0, "y0": y0, "x1": x1, "y1": y1}
        page_height: 页面高度
        threshold: 距离页面底部的阈值（像素）
    Returns:
        bool: 如果表格在页面底部返回True
    """
    # PDF坐标系中，y1越大越靠近页面底部
    distance_from_bottom = page_height - table_bbox["y1"]
    return distance_from_bottom <= threshold

def has_page_number_content_at_bottom(mixed_content, page_height, threshold=50):
    """
    检查页面底部是否有页码内容（排除1-1-16格式的页码）
    """
    if not mixed_content.get("text_blocks"):
        return False
    
    for paragraph in mixed_content["text_blocks"]:
        if paragraph.get("bbox"):
            # 检查是否在页面底部
            distance_from_bottom = page_height - paragraph["bbox"]["y1"]
            if distance_from_bottom <= threshold:
                # 检查是否是页码格式
                if is_page_number_format(paragraph["text"]):
                    return True
    return False

def extract_content_all(file_path, output_txt_path=None):
    """
    提取PDF文档所有内容，按y轴坐标从低到高排列输出到txt文件
    支持跨页表格合并：如果当前页顶部有表格，上一页底部也有表格（非页码），则合并
    
    Args:
        file_path (str): PDF文件路径
        output_txt_path (str, optional): 输出txt文件路径，默认为同名txt文件
        
    Returns:
        str: 输出文件路径
    """
    if output_txt_path is None:
        # 生成默认输出路径
        base_name = os.path.splitext(os.path.basename(file_path))[0]
        output_txt_path = f"./file_output/{base_name}_content_all.txt"
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_txt_path), exist_ok=True)
    
    all_content_items = []  # 存储所有内容项，用于按y轴排序
    previous_page_content = None  # 存储上一页的内容信息
    
    with pdfplumber.open(file_path) as pdf:
        print(f"开始处理PDF文件: {file_path}")
        print(f"总页数: {len(pdf.pages)}")
        
        # 遍历所有页面range(14, 15) enumerate(pdf.pages/pdf.pages[12:16])
        for page_num, page in enumerate(pdf.pages):
            print(f"处理第 {page_num + 1} 页...")
            
            # 获取页面内容和尺寸
            mixed_content = extract_mixed_content(page)
            page_height = page.height
            
            current_page_tables = []
            merged_table_indices = []  # 记录哪些表格被合并了
            
            # 处理表格内容
            if mixed_content["tables"]:
                for table_idx, table_info in enumerate(mixed_content["tables"]):
                    # 使用表格的y0坐标作为排序依据（PDF坐标系：y0越小越在页面顶部）
                    y_position = table_info["bbox"]["y0"]
                    
                    # 检查是否需要与上一页的表格合并
                    should_merge = False
                    merge_target_idx = None
                    
                    # 检查表格合并条件
                    if (previous_page_content and 
                        is_table_at_page_top(table_info["bbox"], page_height) and
                        not has_page_number_content_at_bottom(mixed_content, page_height)):
                        
                        # 查找上一页底部的表格
                        for prev_idx, prev_item in enumerate(previous_page_content):
                            if (prev_item['content_type'] == 'table' and 
                                'table_bbox' in prev_item and
                                is_table_at_page_bottom(prev_item['table_bbox'], page_height)):
                                should_merge = True
                                merge_target_idx = prev_idx
                                break
                    
                    if should_merge and merge_target_idx is not None:
                        print(f"  🔗 合并表格：将第{page_num + 1}页顶部表格合并到第{page_num}页底部表格")
                        
                        # 找到上一页对应的表格在all_content_items中的位置
                        prev_table_item = None
                        prev_table_item_idx = None
                        for idx, item in enumerate(all_content_items):
                            if (item.get('content_type') == 'table' and 
                                item.get('page_num') == page_num and  # 上一页 (page_num是从1开始的)
                                'table_bbox' in item and
                                item['table_bbox'] == previous_page_content[merge_target_idx]['table_bbox']):
                                prev_table_item = item
                                prev_table_item_idx = idx
                                break
                        
                        if prev_table_item and prev_table_item_idx is not None:
                            # 构建当前表格的数据部分（不包含标题）
                            current_table_data = ""
                            for row in table_info['data']:
                                if row:
                                    row_text = '\t'.join(str(cell) if cell else '' for cell in row)
                                    current_table_data += row_text + '\n'
                            
                            # 合并表格内容：将当前表格数据追加到上一页表格内容
                            merged_content = prev_table_item['content'].rstrip() + '\n' + current_table_data + '\n'
                            
                            # 更新all_content_items中上一页表格的内容
                            all_content_items[prev_table_item_idx]['content'] = merged_content
                            
                            # 更新previous_page_content中的内容
                            previous_page_content[merge_target_idx]['content'] = merged_content
                        
                        # 记录这个表格已被合并，不添加到当前页
                        merged_table_indices.append(table_idx)
                        
                    else:
                        # 正常处理表格（不合并）
                        table_text = f"\n[表格 {table_info['table_index']} - 第{page_num + 1}页]\n"
                        table_text += f"({table_info['row_count']}行 x {table_info['col_count']}列)\n"
                        table_text += "-" * 50 + "\n"
                        
                        for row in table_info['data']:
                            if row:
                                row_text = '\t'.join(str(cell) if cell else '' for cell in row)
                                table_text += row_text + '\n'
                        table_text += "\n"
                        
                        table_item = {
                            'page_num': page_num + 1,
                            'y_position': y_position,
                            'content_type': 'table',
                            'content': table_text,
                            'table_bbox': table_info["bbox"]  # 保存bbox信息用于下一页判断
                        }
                        
                        current_page_tables.append(table_item)
                        all_content_items.append(table_item)
            
            # 处理文本块内容（段落）
            current_page_paragraphs = []
            if mixed_content["text_blocks"]:
                for paragraph in mixed_content["text_blocks"]:
                    # 使用段落的y0坐标作为排序依据
                    if paragraph.get("bbox"):
                        y_position = paragraph["bbox"]["y0"]
                    else:
                        # 如果没有bbox信息，使用一个默认值
                        y_position = 0
                    
                    paragraph_item = {
                        'page_num': page_num + 1,
                        'y_position': y_position,
                        'content_type': 'paragraph',
                        'content': paragraph['text'] + '\n\n'
                    }
                    
                    current_page_paragraphs.append(paragraph_item)
                    all_content_items.append(paragraph_item)
            
            # 更新上一页内容信息（用于下一页的表格合并判断）
            previous_page_content = current_page_tables + current_page_paragraphs
        
        # 按页码和y轴位置排序（页码优先，同一页内按y轴从小到大）
        # PDF坐标系中，y值越小越靠近页面顶部
        all_content_items.sort(key=lambda x: (x['page_num'], x['y_position']))
        
        # 写入txt文件
        with open(output_txt_path, 'w', encoding='utf-8') as f:
            f.write(f"PDF文档内容提取结果\n")
            f.write(f"源文件: {file_path}\n")
            f.write(f"提取时间: {__import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"总页数: {len(pdf.pages)}\n")
            f.write(f"内容项总数: {len(all_content_items)}\n")
            f.write("=" * 80 + "\n\n")
            
            current_page = 0
            for item in all_content_items:
                # 如果是新页面，添加页面分隔符
                if item['page_num'] != current_page:
                    if current_page > 0:  # 不是第一页
                        f.write("\n" + "=" * 60 + "\n")
                    f.write(f"第 {item['page_num']} 页\n")
                    f.write("=" * 60 + "\n\n")
                    current_page = item['page_num']
                
                # 写入内容
                f.write(item['content'])
        
        print(f"✅ 内容提取完成!")
        print(f"📄 输出文件: {output_txt_path}")
        print(f"📊 统计信息:")
        print(f"   - 总内容项: {len(all_content_items)}")
        print(f"   - 表格: {sum(1 for item in all_content_items if item['content_type'] == 'table')}")
        print(f"   - 段落: {sum(1 for item in all_content_items if item['content_type'] == 'paragraph')}")
        
        return output_txt_path

# 主函数
def main():
    # 测试 extract_content_all 函数
    try:
        # 假设您有一个测试PDF文件
        pdf_file = "./file_input/test.pdf"
        if os.path.exists(pdf_file):
            print("测试 extract_content_all 函数...")
            output_file = extract_content_all(pdf_file)
            print(f"✅ 测试完成，输出文件: {output_file}")
        else:
            print("⚠️ 测试PDF文件不存在，请将PDF文件放在 ./file_input/test.pdf")
    except Exception as e:
        print(f"❌ 测试失败: {e}")


if __name__ == "__main__":
    main()